﻿#include "modelapplication.h"
#include "ui_modelapplication.h"
#include <QFileDialog>
#include <modelinfodlg.h>
#include <QMessageBox>
#include <QDirIterator>
#include <QDebug>
#include <iostream>


using torch::jit::script::Module;

void ModelApplication::ScanModels()
{
    m_modelsdir = qApp->applicationDirPath().replace("Runtimes","Plugins")+"/AIModel/models/";
    //QMessageBox::information(NULL,QStringLiteral("模型加载"),modelsdir);
    //qDebug() << pluginsdir;
    QDirIterator it(m_modelsdir,QDir::Files|QDir::Dirs|QDir::NoDotAndDotDot);	//遍历所有目录和文件
    while (it.hasNext())//存在
    {
        QString name = it.next();//读取
        QFileInfo info(name);
        if (info.isDir())//判断是目录
        {
            //将目录加入运行环境变量
            continue;
        }
        else
        {
            if (info.suffix() == "pt"){
                QString modelname = info.baseName().toUtf8().data();
                QFile file(info.filePath().replace(".pt",".txt"));
                if(file.open(QIODevice::ReadOnly)){
                    m_modelinfo[modelname] = file.readAll();
                }else{
                    qDebug() << modelname <<":" <<  QStringLiteral("标签文件加载失败");
                    continue;
                }
                ui->listModels->addItem(modelname);
            }
        }
    }
}

ModelApplication::ModelApplication(QWidget *parent) :
    QWidget(parent),
    ui(new Ui::ModelApplication)
{
    ui->setupUi(this);
    this->setWindowTitle(QStringLiteral("模型工具"));
    ScanModels();
    m_type = MEDIA_IMG;
    m_isPath = false;
}

ModelApplication::~ModelApplication()
{
    delete ui;
}

void ModelApplication::on_btnadd_clicked()
{
    QString url = QFileDialog::getOpenFileName(NULL,QStringLiteral("选择模型"),"./","*.pt");
    if(!url.isEmpty()){
        //弹出模型信息编辑框
        ModelInfoDlg dlg;
        if(dlg.exec() == QDialog::Accepted){
            //复制保存模型 模型名称 标签信息
            if(m_modelinfo.contains(dlg.m_modelName)){
                QMessageBox::information(NULL,QStringLiteral("模型导入"),QStringLiteral("模型名称已存在，请重新命名"));
                return;
            }
            QString modelpath = m_modelsdir+dlg.m_modelName+".pt";
            QString labelpath = m_modelsdir+dlg.m_modelName+".txt";
            qDebug() << modelpath;
            if (QFile::exists(modelpath)){
                QFile::remove(modelpath);
            }
            if (QFile::exists(labelpath)){
                QFile::remove(labelpath);
            }
            if(QFile::copy(url,modelpath)){
                qDebug() << QStringLiteral("模型导入成功");
                ui->listModels->addItem(dlg.m_modelName);
                m_modelinfo[dlg.m_modelName] = dlg.m_labelList;
                QFile file(labelpath);
                if(file.open(QIODevice::WriteOnly)){
                    file.write(dlg.m_labelList.toUtf8());
                }
                file.close();
            }else{
                QMessageBox::information(NULL,QStringLiteral("模型导入"),QStringLiteral("模型导入失败"));
            }
        }
    }
}

void ModelApplication::on_btnminus_clicked()
{
}

void ModelApplication::on_listModels_itemDoubleClicked(QListWidgetItem *item)
{
    QMessageBox::information(NULL,QStringLiteral("模型标签"),m_modelinfo[item->text()]);
}

void ModelApplication::on_btnselectimg_clicked()
{
    m_targetPath = QFileDialog::getOpenFileName(NULL,QStringLiteral("打开图片"),"./","PNG Files(*.png);;JPEG Files(*.jpg);;");
    if(!m_targetPath.isEmpty()){
        ui->labelpath->setText(m_targetPath);
        m_isPath = false;
        m_type = MEDIA_IMG;
    }
}

void ModelApplication::on_btnselectdir_clicked()
{
    m_targetPath = QFileDialog::getExistingDirectory(this,QStringLiteral("选择目录"),"./",QFileDialog::ShowDirsOnly);
     if(!m_targetPath.isEmpty()){
         ui->labelpath->setText(m_targetPath);
         m_isPath = false;
         m_type = MEDIA_IMG;
     }
}

void ModelApplication::on_listMethods_itemDoubleClicked(QListWidgetItem *item)
{
    if(item->text() == QStringLiteral("结果展示")){
        if(ui->listModels->selectedItems().size()==0){
            QMessageBox::information(NULL,QStringLiteral("识别结果"),QStringLiteral("请先选择模型"));
            return;
        }
        QString modename = ui->listModels->currentItem()->text();
        if(!m_modelinfo.contains(modename)){
            QMessageBox::information(NULL,QStringLiteral("识别结果"),QStringLiteral("模型不存在:")+modename);
            return;
        }
        auto device = torch::kCPU;
        try{
            Module module = torch::jit::load((m_modelsdir+modename+".pt").toStdString(),device);
        }catch (const c10::Error& e) {
            QMessageBox::information(NULL,QStringLiteral("识别结果"),e.msg().c_str());
            return;
        }

    }
}
